#include "TChain.h"
#include "TTree.h"
#include "TH1F.h"
#include "TH2D.h"
#include "TCanvas.h"
#include "THnSparse.h"

#include "AliAnalysisTask.h"
#include "AliAnalysisManager.h"

#include "AliVEvent.h"
#include "AliESDEvent.h"

#include "AliMCEventHandler.h"
#include "AliMCEvent.h"

#include "AliPhysicsSelection.h"

#include "AliESDtrackCuts.h"
#include "AliStack.h"
#include "TParticle.h"

#include "AliAnalysisTaskMCtrue.h"

#define PT_MIN 0.3  //set range for Pt
#define PT_MAX 1.5  //(if this selection is performed in ESD track loop)

#define ETA_MIN -0.8  //set range for pseudorapidity
#define ETA_MAX 0.8   //(if this selection is performed in ESD track loop)


ClassImp(AliAnalysisTaskMCtrue)

//________________________________________________________________________
AliAnalysisTaskMCtrue::AliAnalysisTaskMCtrue(const char *name) 
  : AliAnalysisTaskSE(name)
, fVxMax(0)
, fVyMax(0)
, fVzMax(0)
, fOutput(0)
, fHistPtPlus(0)
, fHistPtMinus(0)
, fHistPt(0)
, fHistPtPlusAfterPtCut(0)
, fHistPtMinusAfterPtCut(0)
, fHistPtAfterPtCut(0)
, fHistEtaPlus(0)
, fHistEtaMinus(0)
, fHistEta(0)
, fHistEtaPlusAfterPtCut(0)
, fHistEtaMinusAfterPtCut(0)
, fHistEtaAfterPtCut(0)
, fHistEtaAfterCuts(0)
, fHistPtAfterCuts(0)
, fHistEtaPtAfterCuts(0)
, fHistVtxX(0)
, fHistVtxY(0)
, fHistVtxZ(0)
, fHistMClabels(0)
, fHistTrackCharge(0)
, fHistEtaPhiPlus(0)
, fHistEtaPhiMinus(0)
, fHistEtaPhi(0)
, fHistsSparse(0)
, fHistNpEvent(0)
, fHistNmEvent(0)
, fHistPtEtaEvent(0)
, fHistEventCutStats(0)
, fHistTrackCutStats(0)
{
  // Constructor

  // Define input and output slots here
  // Input slot #0 works with a TChain
  DefineInput(0, TChain::Class());
  // Output slot #0 id reserved by the base class for AOD
  // Output slot #1 writes into a TH1 container
  DefineOutput(1, TList::Class());
}

AliAnalysisTaskMCtrue::AliAnalysisTaskMCtrue()
  : AliAnalysisTaskSE("user_task")
, fVxMax(0)
, fVyMax(0)
, fVzMax(0)
, fOutput(0)
, fHistPtPlus(0)
, fHistPtMinus(0)
, fHistPt(0)
, fHistPtPlusAfterPtCut(0)
, fHistPtMinusAfterPtCut(0)
, fHistPtAfterPtCut(0)
, fHistEtaPlus(0)
, fHistEtaMinus(0)
, fHistEta(0)
, fHistEtaPlusAfterPtCut(0)
, fHistEtaMinusAfterPtCut(0)
, fHistEtaAfterPtCut(0)
, fHistEtaAfterCuts(0)
, fHistPtAfterCuts(0)
, fHistEtaPtAfterCuts(0)
, fHistVtxX(0)
, fHistVtxY(0)
, fHistVtxZ(0)
, fHistMClabels(0)
, fHistTrackCharge(0)
, fHistEtaPhiPlus(0)
, fHistEtaPhiMinus(0)
, fHistEtaPhi(0)
, fHistsSparse(0)
, fHistNpEvent(0)
, fHistNmEvent(0)
, fHistPtEtaEvent(0)
, fHistEventCutStats(0)
, fHistTrackCutStats(0)
{
  // Constructor

  // Define input and output slots here
  // Input slot #0 works with a TChain
  DefineInput(0, TChain::Class());
  // Output slot #0 id reserved by the base class for AOD
  // Output slot #1 writes into a TH1 container
  DefineOutput(1, TList::Class());
}

AliAnalysisTaskMCtrue::~AliAnalysisTaskMCtrue()
{
  // Destructor
}

//________________________________________________________________________
void AliAnalysisTaskMCtrue::UserCreateOutputObjects()
{
  // Create histograms
  // Called once
  //
  // For histograms description see .h file

  fOutput = new TList();

  fHistPtPlus = new TH1F("fPtPlus", "+ Particles Vs Pt", 1000, 0, 10 );
  fHistPtPlus->GetXaxis()->SetTitle("GeV");
  fHistPtPlus->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistPtPlus);

  fHistPtMinus = new TH1F("fPtMinus", "- Particles Vs Pt", 1000, 0, 10 );
  fHistPtMinus->GetXaxis()->SetTitle("GeV");
  fHistPtMinus->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistPtMinus);

  fHistPt = new TH1F("fPt", "ch Particles Vs Pt", 1000, 0, 10 );
  fHistPt->GetXaxis()->SetTitle("GeV");
  fHistPt->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistPt);

  fHistPtPlusAfterPtCut = new TH1F("fPtPlusAfterPtCut", "+ Particles Vs Pt after Pt cut", 300, 0, 3 );
  fHistPtPlusAfterPtCut->GetXaxis()->SetTitle("GeV");
  fHistPtPlusAfterPtCut->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistPtPlusAfterPtCut);

  fHistPtMinusAfterPtCut = new TH1F("fPtMinusAfterPtCut", "- Particles Vs Pt after Pt cut", 300, 0, 3 );
  fHistPtMinusAfterPtCut->GetXaxis()->SetTitle("GeV");
  fHistPtMinusAfterPtCut->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistPtMinusAfterPtCut);

  fHistPtAfterPtCut = new TH1F("fPtAfterPtCut", "ch Particles Vs Pt after Pt cut", 300, 0, 3 );
  fHistPtAfterPtCut->GetXaxis()->SetTitle("GeV");
  fHistPtAfterPtCut->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistPtAfterPtCut);

  fHistEtaPlus = new TH1F("fEtaPlus", "+ particles vs Eta", 400, -2, 2 );
  fHistEtaPlus->GetXaxis()->SetTitle("#eta");
  fHistEtaPlus->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistEtaPlus);

  fHistEtaMinus = new TH1F("fEtaMinus", "- particles vs Eta", 400, -2, 2 );
  fHistEtaMinus->GetXaxis()->SetTitle("#eta");
  fHistEtaMinus->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistEtaMinus);

  fHistEta = new TH1F("fEta", "ch particles vs Eta", 400, -2, 2 );
  fHistEta->GetXaxis()->SetTitle("#eta");
  fHistEta->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistEta);

  fHistEtaPlusAfterPtCut = new TH1F("fEtaPlusAfterPtCut", "+ particles vs Eta after Pt Cut", 400, -2, 2 );
  fHistEtaPlusAfterPtCut->GetXaxis()->SetTitle("#eta");
  fHistEtaPlusAfterPtCut->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistEtaPlusAfterPtCut);

  fHistEtaMinusAfterPtCut = new TH1F("fEtaMinusAfterPtCut", "- particles vs Eta after Pt Cut", 400, -2, 2 );
  fHistEtaMinusAfterPtCut->GetXaxis()->SetTitle("#eta");
  fHistEtaMinusAfterPtCut->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistEtaMinusAfterPtCut);

  fHistEtaAfterPtCut = new TH1F("fEtaAfterPtCut", "ch particles vs Eta after Pt Cut", 400, -2, 2 );
  fHistEtaAfterPtCut->GetXaxis()->SetTitle("#eta");
  fHistEtaAfterPtCut->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistEtaAfterPtCut);

  fHistEtaAfterCuts = new TH1F("fEtaAfterCuts", "ch particles vs Eta after all cuts", 400, -2, 2 );
  fHistEtaAfterCuts->GetXaxis()->SetTitle("#eta");
  fHistEtaAfterCuts->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistEtaAfterCuts);

  fHistPtAfterCuts = new TH1F("fPtAfterCuts", "ch Particles Vs Pt after all cuts", 300, 0, 3 );
  fHistPtAfterCuts->GetXaxis()->SetTitle("GeV");
  fHistPtAfterCuts->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistPtAfterCuts);

  fHistEtaPtAfterCuts = new TH2F("fEtaPtAfterCuts", "ch Particles Vs Eta and Pt after all cuts", 200, -1, 1, 200, 0, 2);
  fHistEtaPtAfterCuts->GetXaxis()->SetTitle("#eta");
  fHistEtaPtAfterCuts->GetYaxis()->SetTitle("GeV");
  fOutput->Add(fHistEtaPtAfterCuts);

  fHistVtxX = new TH1F("fVtxX", "Vertex x coordinate distribution", 100, -0.5, 0.5 );
  fHistVtxX->GetXaxis()->SetTitle("Vx (cm)");
  fHistVtxX->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistVtxX);

  fHistVtxY = new TH1F("fVtxY", "Vertex y coordinate distribution", 100, -0.5, 0.5 );
  fHistVtxY->GetXaxis()->SetTitle("Vy (cm)");
  fHistVtxY->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistVtxY);

  fHistVtxZ = new TH1F("fVtxZ", "Vertex z coordinate distribution", 400, -20, 20 );
  fHistVtxZ->GetXaxis()->SetTitle("Vz (cm)");
  fHistVtxZ->GetYaxis()->SetTitle("N entries");
  fOutput->Add(fHistVtxZ);

  fHistMClabels = new TH1I("fHistMClabels","MC label;label;Entries",221,-10.5,210.5);
  fOutput->Add(fHistMClabels);
  
  fHistTrackCharge = new TH1D("fHistTracksCharge","tracks charge;charge;Entries",11,-5.5,5.5);
  fOutput->Add(fHistTrackCharge);

  fHistEtaPhiPlus = new TH2F("fEtaPhiPlus", "+ particles vs EtaPhi after all cuts", 40, -1, 1, 71, -0.515, 6.798 );//2pi!! 6.283/61=.103
  fHistEtaPhiPlus->GetXaxis()->SetTitle("#eta");
  fHistEtaPhiPlus->GetYaxis()->SetTitle("#phi");
  fOutput->Add(fHistEtaPhiPlus);

  fHistEtaPhiMinus = new TH2F("fEtaPhiMinus", "- particles vs EtaPhi after all cuts",40,-1,1,71,-0.515,6.798 );//2pi!!
  fHistEtaPhiMinus->GetXaxis()->SetTitle("#eta");
  fHistEtaPhiMinus->GetYaxis()->SetTitle("#phi");
  fOutput->Add(fHistEtaPhiMinus);

  fHistEtaPhi = new TH2F("fEtaPhi", "N particles vs EtaPhi after all cuts", 40, -1, 1, 71, -0.515, 6.798 );//2pi!!
  fHistEtaPhi->GetXaxis()->SetTitle("#eta");
  fHistEtaPhi->GetYaxis()->SetTitle("#phi");
  fOutput->Add(fHistEtaPhi);

  //for pp:  {Np_B,Np_F,Nm_B,Nm_F,Nch_B,Nch_F,Q_B, Q_F,PtMean_B,PtMean_F,Q_F*Q_B,Q_F*Q_F }
  Int_t bins[12] = {  50,  50,  50,  50, 100, 100,  50,   50,  140,  140,   500,   400};//number of bins
  Double_t min[12]={-0.5,-0.5,-0.5,-0.5,-0.5,-0.5,-25.5,-25.5,0.195,0.195, -250.5, -0.5};//min values
  Double_t max[12]={49.5,49.5,49.5,49.5,99.5,99.5, 24.5, 24.5,1.595,1.595,  249.5, 399.5};//max values

  fHistsSparse = new THnSparseF*[16];
  for (Int_t i=0; i<=9; i++)
  {
  char str[]="fSparse0";
  str[7]+=i;
  fHistsSparse[i] = new THnSparseF(str,"Sparse hist",12,bins,min,max);
  fOutput->Add(fHistsSparse[i]);
  }

  for (Int_t i=10; i<=15; i++)
  {
  char str1[]="fSparse10";
  str1[8]+=(i-10);
  fHistsSparse[i] = new THnSparseF(str1,"Sparse hist",12,bins,min,max);
  fOutput->Add(fHistsSparse[i]);
  }

  fHistNpEvent = new TH1F("fNpEvent", "Nplus in one event vs Eta", 16, -0.8, 0.8 );

  fHistNmEvent = new TH1F("fNmEvent", "Nminus in one event vs Eta", 16, -0.8, 0.8 );

  fHistPtEtaEvent = new TH2F("f2PtEtaEvent", "Pt vs Eta in one event", 140, 0.195, 1.595, 16, -0.8, 0.8);

  //Event level QA
  Int_t nStatBins = 8;
  fHistEventCutStats = new TH1D("fHistEventCutStats","Event statistics;;N_{events}", nStatBins,-0.5,nStatBins-0.5);
  TString gEventCutBinNames[8] = {"Total","No trigger","Centrality","No vertex","Bad vertex position","fESD not available","LowMult cut","Analyzed"};
  
  for(Int_t i = 1; i <= nStatBins; i++){
  fHistEventCutStats->GetXaxis()->SetBinLabel(i,gEventCutBinNames[i-1].Data());
  }
  fOutput->Add(fHistEventCutStats);

  //Track level QA
  fHistTrackCutStats = new TH1D("fHistTrackCutStats","Event statistics;N_{events}", 9,-0.5,8.5);
  TString gTrackCutBinNames[9] = {"Total","Label < 0","Not primary","Charge 0","HighPt cut","LowPt cut","High Eta cut","Low Eta Cut","Analyzed"};
  
  for(Int_t i = 1; i <= 9; i++){
  fHistTrackCutStats->GetXaxis()->SetBinLabel(i,gTrackCutBinNames[i-1].Data());
  }
  fOutput->Add(fHistTrackCutStats);

  //fOutput->Add(fTrackCuts);
}

//________________________________________________________________________
void AliAnalysisTaskMCtrue::UserExec(Option_t *) 
{

  // Main loop
  // Called for each event

  //if ( fHaveMC )
  //	UserExecMC();

  AliESDEvent *fESD = dynamic_cast<AliESDEvent*>(InputEvent());

  fHistEventCutStats->Fill(0); //total number of events

  if ( !fESD ) 
  {
	//Printf("ERROR: fESD not available\n");
        fHistEventCutStats->Fill(5);
	return;
  }

    //look at the kinematics
    AliMCEvent *eventMC = MCEvent();
    AliStack *stack = eventMC->Stack();


    //Trigger
    Bool_t lTrigger = kTRUE;
    if( fESD )
    {
        lTrigger = ((AliInputEventHandler*)(AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler()))->IsEventSelected();
        if( !lTrigger )
        {
            fHistEventCutStats->Fill(1);
            PostData(1, fOutput);
            return;
        }
    }

   // Vertex present

    const AliESDVertex *vertex = fESD->GetPrimaryVertex();

     Bool_t lVertexPresent = ( (vertex) && (vertex->GetNContributors() > 0 ) && ( vertex->GetZRes() != 0 ) );
     if( !lVertexPresent ) 
    {
            fHistEventCutStats->Fill(3); //no vertex
            PostData(1, fOutput);
            return;
    }
     

  // Vertex in range
     Bool_t lVertexAcceptable = (TMath::Abs(vertex->GetXv()) < fVxMax) && (TMath::Abs(vertex->GetYv()) < fVyMax);
     if( lVertexAcceptable )
   {
     if( fVzMax>0 )
     lVertexAcceptable = (TMath::Abs(vertex->GetZv()) < fVzMax);
     else
     lVertexAcceptable = (TMath::Abs(vertex->GetZv()) > -fVzMax);
   }

    if(!lVertexAcceptable)
   {
            fHistEventCutStats->Fill(4); //bad vertex
            PostData(1, fOutput);
            return;
   }

    Double_t lVtxX = vertex->GetXv();
    fHistVtxX->Fill(lVtxX);//x coordinate distribution of vertex

    Double_t lVtxY = vertex->GetYv();
    fHistVtxY->Fill(lVtxY);//y coordinate distribution of vertex

    Double_t lVtxZ = vertex->GetZv();
    fHistVtxZ->Fill(lVtxZ);//z coordinate distribution of vertex

  // user trigger to be added (if needed)
  // user vertex cut (if needed)

  fHistNpEvent->Reset();
  fHistNmEvent->Reset();
  fHistPtEtaEvent->Reset();


  //Track loop
  for (Int_t iTracks = 0; iTracks < eventMC->GetNumberOfTracks(); iTracks++) 
  {
    AliVParticle* track = eventMC->GetTrack(iTracks);

    fHistTrackCutStats->Fill(0); //total number of tracks

    Int_t label = track->GetLabel();
    fHistMClabels->Fill(label);


      //by IA: by Peter Hristov - it's ghost tracks
      if ( label < 0 )
     {

     fHistTrackCutStats->Fill(1);
     continue;
     }

     if(! stack->IsPhysicalPrimary(label))
     {

       fHistTrackCutStats->Fill(2);
       continue;
     }


     Short_t lTrackCharge = track->Charge()/3;//track charge in correct values |e|


     
     if (lTrackCharge == 0)
     {

       fHistTrackCutStats->Fill(3);
       continue;
     }

     if ( !track ) 
    {
	//printf("ERROR: Could not receive track %d\n", iTracks);
	continue;
    }

	Double_t lTrackPt = track->Pt();
	Double_t lTrackEta = track->Eta();




    if (lTrackCharge>0)
    {
    fHistPtPlus->Fill(lTrackPt);
    fHistEtaPlus->Fill(lTrackEta);
    }

    if (lTrackCharge<0)
    {
    fHistPtMinus->Fill(lTrackPt);
    fHistEtaMinus->Fill(lTrackEta);
    }

    fHistPt->Fill(lTrackPt);
    fHistEta->Fill(lTrackEta);

        if ( lTrackPt > PT_MAX)
        {fHistTrackCutStats->Fill(4);}
        if ( lTrackPt < PT_MIN)
        {fHistTrackCutStats->Fill(5);}

	//cut on Pt
	if( lTrackPt < PT_MIN || lTrackPt > PT_MAX )
	continue;

       if (lTrackCharge > 0)
       {
       fHistPtPlusAfterPtCut->Fill(lTrackPt);
       fHistEtaPlusAfterPtCut->Fill(lTrackEta);
       }

       if (lTrackCharge < 0)
       {
       fHistPtMinusAfterPtCut->Fill(lTrackPt);
       fHistEtaMinusAfterPtCut->Fill(lTrackEta);
       }

    fHistPtAfterPtCut->Fill(lTrackPt);
    fHistEtaAfterPtCut->Fill(lTrackEta);

        if ( lTrackEta > ETA_MAX)
        {fHistTrackCutStats->Fill(6);}
        if ( lTrackEta < ETA_MIN)
        {fHistTrackCutStats->Fill(7);}

        //cut on Eta
	if( lTrackEta < ETA_MIN || lTrackEta > ETA_MAX )
	continue;

       fHistEtaAfterCuts->Fill(lTrackEta);
       fHistPtAfterCuts->Fill(lTrackPt);
       fHistEtaPtAfterCuts->Fill(lTrackEta, lTrackPt);

        Double_t lTrackPhi = track->Phi();

        if (lTrackCharge > 0) {fHistNpEvent->Fill(lTrackEta);}
        if (lTrackCharge < 0) {fHistNmEvent->Fill(lTrackEta);}

        fHistEtaPhi->Fill( lTrackEta, lTrackPhi);
        if (lTrackCharge > 0) { fHistEtaPhiPlus->Fill( lTrackEta, lTrackPhi);}
        if (lTrackCharge < 0) { fHistEtaPhiMinus->Fill( lTrackEta, lTrackPhi);}

        fHistPtEtaEvent->Fill(lTrackPt,lTrackEta);
        fHistTrackCharge->Fill(lTrackCharge);
        fHistTrackCutStats->Fill(8);
  }

   //arrays of bin numbers for 16 eta windows
   Int_t downB[16]={ 7, 6, 5, 4, 3, 2, 1, 5, 4, 3, 2, 1, 3, 2, 1, 1};
   Int_t upB[16] = { 8, 7, 6, 5, 4, 3, 2, 8, 7, 6, 5, 4, 8, 7, 6, 8};
   Int_t downF[16]={ 9,10,11,12,13,14,15, 9,10,11,12,13, 9,10,11, 9};
   Int_t upF[16] = {10,11,12,13,14,15,16,12,13,14,15,16,14,15,16,16};

   for (Int_t i=0;i<=15;i++)
   {
   Double_t Np_B = fHistNpEvent->Integral(downB[i],upB[i]);//Np backward
   Double_t Np_F = fHistNpEvent->Integral(downF[i],upF[i]);//Np forward
   Double_t Nm_B = fHistNmEvent->Integral(downB[i],upB[i]);//Nm backward
   Double_t Nm_F = fHistNmEvent->Integral(downF[i],upF[i]);//Nm forward
   Double_t Nch_B = Np_B + Nm_B;
   Double_t Nch_F = Np_F + Nm_F;
   Double_t Q_B = Np_B - Nm_B;
   Double_t Q_F = Np_F - Nm_F;
   Double_t PtMean_B = fHistPtEtaEvent->ProjectionX("name2",downB[i],upB[i])->GetMean();//backward
   Double_t PtMean_F = fHistPtEtaEvent->ProjectionX("name",downF[i],upF[i])->GetMean();//forward
   Double_t x[12]={Np_B, Np_F, Nm_B, Nm_F, Nch_B, Nch_F, Q_B, Q_F, PtMean_B, PtMean_F, Q_B*Q_F, Q_F*Q_F};
   fHistsSparse[i]->Fill(x);
   }


   //fHistCent->Fill(lCentrality);
   fHistEventCutStats->Fill(7); //analyzed

   PostData(1, fOutput);
}



//________________________________________________________________________
void AliAnalysisTaskMCtrue::Terminate(Option_t *) 
{
  // Draw result to the screen
  // Called once at the end of the query

 
  fOutput = dynamic_cast<TList*> (GetOutputData(1));
  if (!fOutput) {
    printf("ERROR: Output list not available\n");
    return;
  }

}
